package vip.liux.backend.infrastructure.config;

import org.aopalliance.intercept.MethodInterceptor;
import org.aopalliance.intercept.MethodInvocation;
import org.apache.commons.lang3.StringUtils;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.slf4j.MDC;
import org.springframework.amqp.core.Message;
import org.springframework.amqp.core.MessageListener;
import org.springframework.amqp.rabbit.config.ContainerCustomizer;
import org.springframework.amqp.rabbit.core.RabbitTemplate;
import org.springframework.amqp.rabbit.listener.MessageListenerContainer;
import org.springframework.aop.framework.ProxyFactory;
import org.springframework.boot.autoconfigure.amqp.RabbitTemplateCustomizer;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.UUID;

import static vip.liux.core.constant.GlobalConstant.MDC_SPAN_ID;
import static vip.liux.core.constant.GlobalConstant.MDC_TRACE_ID;

@Configuration(proxyBeanMethods = false)
public class RabbitConfiguration implements RabbitTemplateCustomizer {

    private static final Logger log = LoggerFactory.getLogger("RabbitConfiguration");


    @Bean
    public ContainerCustomizer<?> containerCustomizer() {
        return (ContainerCustomizer<MessageListenerContainer>) container -> {
            Object listener = container.getMessageListener();
            // Ensure the listener is a public class
            if (listener instanceof MessageListener messageListener) {
                ProxyFactory pf = new ProxyFactory();
                pf.setTarget(messageListener);
                pf.addAdvice(new MethodInterceptor() {
                    @Override
                    public @Nullable Object invoke(@NotNull MethodInvocation invocation) throws Throwable {
                        if ("onMessage".equals(invocation.getMethod().getName())) {
                            Message message = (Message) invocation.getArguments()[0];
                            String spanId = message.getMessageProperties().getHeader(MDC_SPAN_ID);
                            if (StringUtils.isBlank(spanId)) {
                                spanId = UUID.randomUUID().toString();
                            }
                            MDC.put(MDC_SPAN_ID, spanId);
                            String traceId = message.getMessageProperties().getHeader(MDC_TRACE_ID);
                            if (StringUtils.isBlank(traceId)) {
                                spanId = UUID.randomUUID().toString();
                            }
                            MDC.put(MDC_TRACE_ID, traceId);
                            log.debug("Get spanId from message: {},traceId from message: {}", spanId, traceId);
                        }
                        return invocation.proceed();
                    }
                });
                MessageListener proxy = (MessageListener) pf.getProxy();
                container.setupMessageListener(proxy);
            }
        };
    }

    /**
     * Callback to customize a {@link RabbitTemplate} instance.
     *
     * @param rabbitTemplate the rabbitTemplate to customize
     */
    @Override
    public void customize(RabbitTemplate rabbitTemplate) {
        rabbitTemplate.setBeforePublishPostProcessors(message -> {
            String spanId = MDC.get(MDC_SPAN_ID);
            if (StringUtils.isBlank(spanId)) {
                spanId = UUID.randomUUID().toString();
            }
            message.getMessageProperties().setHeader(MDC_SPAN_ID, spanId);

            String traceId = MDC.get(MDC_TRACE_ID);
            if (StringUtils.isBlank(traceId)) {
                traceId = UUID.randomUUID().toString();
            }
            message.getMessageProperties().setHeader(MDC_TRACE_ID, traceId);
            log.info("Set spanId to message: {},traceId to message: {}", spanId, traceId);

            return message;
        });
    }
}
